import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from monai import transforms
from monai.data import MetaTensor
import SimpleITK as sitk
import json
import os
from random import random, sample, seed


def datafold_read(datalist, basedir, fold=0, key="training"):
    with open(datalist) as f:
        json_data = json.load(f)

    json_data = json_data[key]

    for d in json_data:
        for k, v in d.items():
            if isinstance(d[k], list):
                d[k] = [os.path.join(basedir, iv) for iv in d[k]]
            elif isinstance(d[k], str):
                d[k] = os.path.join(basedir, d[k]) if len(d[k]) > 0 else d[k]

    tr = []
    val = []
    for d in json_data:
        if "fold" in d and d["fold"] == fold:
            val.append(d)
        else:
            tr.append(d)

    return tr, val


class BasicDataset(Dataset):
    def __init__(self, data_list, args, phase='train'):
        super(BasicDataset, self).__init__()
        self.data_list = data_list
        self.args = args
        self.train_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.phase = phase

    def __len__(self):
        return len(self.data_list)
    
    def __getitem__(self, item):
        data = self.load_image(self.data_list[item])
        if self.phase == 'train':
            data = self.train_transform(data)
        elif self.phase == 'val':
            data = self.val_transform(data)
        elif self.phase == 'test':
            data = self.test_transform(data)

        return data

    def load_image(self, file_dic):
        image_path = file_dic['image']
        label_path = file_dic['label']
        
        img = sitk.ReadImage(image_path)
        label = sitk.ReadImage(label_path)

        img = sitk.GetArrayFromImage(img).astype(np.float_)
        label = sitk.GetArrayFromImage(label).astype(np.float_)

        img = np.moveaxis(img, (0, 1, 2), (2, 1, 0))
        label = np.moveaxis(label, (0, 1, 2), (2, 1, 0))
        
        return{
            'image': img,
            'label': label,
            'path': image_path
        }


class MSDHeartDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDHeartDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z),
                                       mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandCropByPosNegLabeld(
                    keys=["image", "label"],
                    label_key="label",
                    spatial_size=(args.roi_x, args.roi_y, args.roi_z),
                    pos=2,
                    neg=1,
                    num_samples=1,
                    image_key="image",
                    image_threshold=0,
                ),
                
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
                transforms.RandRotate90d(keys=["image", "label"], prob=0.1, max_k=3),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.1),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.1),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase


class MSDLiverDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDLiverDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z),
                                       mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandCropByPosNegLabeld(
                    keys=["image", "label"],
                    label_key="label",
                    spatial_size=(args.roi_x, args.roi_y, args.roi_z),
                    pos=1,
                    neg=1,
                    num_samples=1,
                    image_key="image",
                    image_threshold=0,
                ),
                transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=2),
                transforms.RandRotate90d(keys=["image", "label"], prob=0.2, max_k=3),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.1),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.1),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase


class MSDHippocampusDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDHippocampusDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z), mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandSpatialCropd(
                    keys=["image", "label"], roi_size=[args.roi_x, args.roi_y, args.roi_z], random_size=False
                ),
                transforms.RandFlipd(keys=["image", "label"], prob=0.1, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.1, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.1, spatial_axis=2),
                transforms.RandRotate90d(keys=["image", "label"], prob=0.1, max_k=3),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.1),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.1),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase


class MSDProstateDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDProstateDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["label"], channel_dim='no_channel'),
                transforms.EnsureChannelFirstd(keys=["image"], channel_dim=-1),
                
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),

                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z), mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandSpatialCropd(
                    keys=["image", "label"], roi_size=[args.roi_x, args.roi_y, args.roi_z], random_size=False
                ),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
                transforms.RandRotate90d(keys=["image", "label"], prob=0.5, max_k=3),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["label"], channel_dim='no_channel'),
                transforms.EnsureChannelFirstd(keys=["image"], channel_dim=-1),

                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase

    def __getitem__(self, item):
        data = self.load_image(self.data_list[item])
        if self.phase == 'train':
            data = self.train_transform(data)
        elif self.phase == 'val':
            data = self.val_transform(data)
        elif self.phase == 'test':
            data = self.test_transform(data)

        return data

    def load_image(self, file_dic):
        image_path = file_dic['image']

        label_path = file_dic['label']
        return {
            'image': image_path,
            'label': label_path,
            'path':  image_path
        }


class MSDLungDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDLungDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", 'label'], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z), mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandCropByPosNegLabeld(
                    keys=["image", "label"],
                    label_key="label",
                    spatial_size=(args.roi_x, args.roi_y, args.roi_z),
                    pos=2,
                    neg=1,
                    num_samples=1,
                    image_key="image",
                    image_threshold=0,
                ),

                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
                transforms.RandRotate90d(keys=["image", "label"], prob=0.3, max_k=3),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.1),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.1),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", 'label'], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase


class MSDPancreasDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDPancreasDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z), mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandCropByPosNegLabeld(
                    keys=["image", "label"],
                    label_key="label",
                    spatial_size=(args.roi_x, args.roi_y, args.roi_z),
                    pos=1,
                    neg=1,
                    num_samples=1,
                    image_key="image",
                    image_threshold=0,
                ),

                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
                transforms.RandRotate90d(keys=["image", "label"], prob=0.25, max_k=3),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        
        self.val_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z), mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase


class MSDHepaticVesselDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDHepaticVesselDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z), mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandSpatialCropd(
                    keys=["image", "label"], roi_size=[args.roi_x, args.roi_y, args.roi_z], random_size=False
                ),

                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
                transforms.RandRotate90d(keys=["image", "label"], prob=0.25, max_k=3),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        
        self.val_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase


class MSDSpleenDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDSpleenDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z), mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandSpatialCropd(
                    keys=["image", "label"], roi_size=[args.roi_x, args.roi_y, args.roi_z], random_size=False
                ),

                transforms.RandFlipd(keys=["image", "label"], prob=0.15, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.15, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.15, spatial_axis=2),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.1),
                transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=0.1),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        
        self.val_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.Spacingd(
                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z),
                    mode=("bilinear", "nearest")
                ),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase


class MSDColonDataset(BasicDataset):
    def __init__(self, data_list, args, phase='train'):
        super(MSDColonDataset, self).__init__(data_list, args, phase='train')
        self.data_list = data_list

        self.train_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.SpatialPadd(["image", "label"], spatial_size=(args.roi_x, args.roi_y, args.roi_z), mode='symmetric'),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image", margin=1),
                transforms.RandCropByPosNegLabeld(
                    keys=["image", "label"],
                    label_key="label",
                    spatial_size=(args.roi_x, args.roi_y, args.roi_z),
                    pos=1,
                    neg=1,
                    num_samples=1,
                    image_key="image",
                    image_threshold=0,
                ),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
                transforms.RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
                transforms.RandRotate90d(keys=["image", "label"], prob=0.25, max_k=3),
                transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )
        self.val_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.test_transform = transforms.Compose(
            [
                # transforms.LoadImaged(keys=["image", "label"]),
                transforms.EnsureChannelFirstd(keys=["image", "label"], channel_dim='no_channel'),
                transforms.ScaleIntensityRanged(
                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
                ),
                transforms.ToTensord(keys=["image", "label"]),
            ]
        )

        self.phase = phase


def get_loader(datasets,
               datalist_json,
               data_dir,
               fold,
               batch_size,
               num_works,
               args=None,
               phase=None):
    
    train_files, validation_files = datafold_read(datalist=datalist_json,
                                                  basedir=data_dir,
                                                  fold=fold)
    seed(12)
    sample_num = np.ceil(len(train_files) * args.frac).astype(np.int_)
    train_files = sample(train_files, sample_num)

    if phase == 'test':
        test_datasets = eval(datasets)(data_list=validation_files, phase='test', args=args)
        test_dataloader = DataLoader(test_datasets,
                                 batch_size=1,
                                 num_workers=num_works,
                                 pin_memory=True,
                                 shuffle=False)
        return test_dataloader
    else:
        train_datasets = eval(datasets)(data_list=train_files, phase='train', args=args)
        val_datasets = eval(datasets)(data_list=validation_files, phase='val', args=args)

        train_dataloader = DataLoader(train_datasets,
                                    batch_size=batch_size,
                                    num_workers=num_works,
                                    pin_memory=True,
                                    shuffle=True)
        # 数据集图像尺寸不一致，推断时使用 滑动窗口推断 输入全尺寸，batch_size 设置为 1 防止无法堆叠
        val_dataloader = DataLoader(val_datasets,
                                    batch_size=1,
                                    num_workers=num_works,
                                    pin_memory=True,
                                    shuffle=False)
        return train_dataloader, val_dataloader

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description="Swin UNETR segmentation pipeline for BRATS Challenge")
    parser.add_argument("--a_min", default=-175.0, type=float, help="a_min in ScaleIntensityRanged")
    parser.add_argument("--a_max", default=250.0, type=float, help="a_max in ScaleIntensityRanged")
    parser.add_argument("--b_min", default=0.0, type=float, help="b_min in ScaleIntensityRanged")
    parser.add_argument("--b_max", default=1.0, type=float, help="b_max in ScaleIntensityRanged")
    parser.add_argument("--space_x", default=1.5, type=float, help="spacing in x direction")
    parser.add_argument("--space_y", default=1.5, type=float, help="spacing in y direction")
    parser.add_argument("--space_z", default=2.0, type=float, help="spacing in z direction")
    parser.add_argument("--roi_x", default=96, type=int, help="roi size in x direction")
    parser.add_argument("--roi_y", default=96, type=int, help="roi size in y direction")
    parser.add_argument("--roi_z", default=96, type=int, help="roi size in z direction")
    parser.add_argument("--dropout_rate", default=0.0, type=float, help="dropout rate")
    parser.add_argument("--frac", default=1.0, type=float, help="dropout rate")

    args = parser.parse_args()

    datasets = 'MSDProstateDataset'
    data_dir = '/home/qlc/raid/dataset/MSD/Task05_Prostate'
    datalist_json = '/home/qlc/model/GMM/MSD/jsons/msd05_folds.json'

    train_files, validation_files = datafold_read(datalist=datalist_json,
                                                  basedir=data_dir,
                                                  fold=0)
    
    # train_transform = transforms.LoadImaged(keys=["image", "label"])
    
    # for path in train_files:
    #     out = train_transform(path)
    #     print(out['image'].shape)
    #     print(out['label'].shape)
    #     break

    train_loader, val_loader = get_loader(datasets, datalist_json, data_dir,
                                          fold=0, batch_size=1, num_works=8, args=args)

    for i, data in enumerate(train_loader):
        print(data['image'].shape)
        print(data['label'].shape)

